from collections.abc import Generator
from typing import Any
import os
import re

from playwright.sync_api import sync_playwright, Browser, APIRequestContext
import pytest
import paramiko

from utils.client import KubeClient
from common.auth import check_storage_file_exists, login
from common.env import load_pw_config, load_ssh_config


@pytest.fixture(scope='session')
def browser_context_args() -> dict[str, Any]:
    pw_config = load_pw_config()
    context_config = pw_config['context']
    return {
        'base_url': context_config['base_url'],
        'storage_state': context_config['storage_state'],
        'ignore_https_errors': True,
    }


@pytest.fixture(scope='session', autouse=True)
def login_and_store_state(browser_context_args) -> None:
    pw_config = load_pw_config()
    context_config = pw_config['context']
    check_storage_file_exists(context_config['storage_state'])
    with sync_playwright() as p:
        browser = p.chromium.launch()
        context = browser.new_context(**browser_context_args)
        login(context, pw_config['credentials'])
        context.storage_state(path=context_config['storage_state'])
        browser.close()


@pytest.fixture(scope='session')
def api_context(browser: Browser, browser_context_args: dict) -> Generator[APIRequestContext, None, None]:
    context = browser.new_context(**browser_context_args)
    api_context = context.request
    yield api_context
    context.close()


@pytest.fixture(scope="session")
def ssh_client() -> Generator[paramiko.SSHClient, None, None]:
    ssh_config = load_ssh_config()
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(**ssh_config)
    yield ssh
    ssh.close()


@pytest.fixture(scope="session")
def kube_client(ssh_client: paramiko.SSHClient) -> Generator[KubeClient, None, None]:
    yield KubeClient(ssh_client)


@pytest.fixture(scope="session")
def numa_stat(ssh_client: paramiko.SSHClient) -> dict:
    try:
        # 获取NUMA节点的数量
        stdin, stdout, stderr = ssh_client.exec_command(
            "ls /sys/devices/system/node/ | grep '^node' | wc -l"
        )
        numa_num_str = stdout.read().decode('utf-8').strip()
        error = stderr.read().decode('utf-8')

        if error:
            raise Exception(f"获取NUMA节点失败: {error}")

        numa_num = int(numa_num_str) if numa_num_str.isdigit() else 1

        if numa_num == 0:
            numa_num = 1

        # 获取每个NUMA节点的CPU数量
        stdin, stdout, stderr = ssh_client.exec_command(
            "ls /sys/devices/system/node/node0/ | grep -c '^cpu[0-9]\+'"
        )
        cpu_count_str = stdout.read().decode('utf-8').strip()
        error = stderr.read().decode('utf-8')

        if error:
            raise Exception(f"获取CPU数量失败: {error}")

        cpu_num_per_numa = int(cpu_count_str) if cpu_count_str.isdigit() else 1

        return {
            "numa": numa_num,
            "cpus": numa_num * cpu_num_per_numa,
            "per_numa": cpu_num_per_numa,
        }

    except Exception as e:
        raise Exception(f"获取NUMA节点失败: {str(e)}") from e
